import os.path

import cv2
import numpy as np


class CvImageDistance:

    def from_file_get_distanct(self, tag_img_path, background_img_path):
        '''
        根据文件进行识别
        :param tag_img_path: 滑块图片的文件路径
        :param background_img_path: 背景图片的文件路径
        :return:
        '''
        # 使用支持中文路径的方式读取图片
        def imread_chinese(filepath, flags=cv2.IMREAD_COLOR):
            """读取包含中文路径的图片"""
            img_array = np.fromfile(filepath, dtype=np.uint8)
            return cv2.imdecode(img_array, flags)
        
        target = imread_chinese(tag_img_path)
        # 读取到两个图片，进行灰值化处理
        template = imread_chinese(background_img_path, cv2.IMREAD_GRAYSCALE)
        
        # 检查图片是否读取成功
        if target is None or template is None:
            raise ValueError("无法读取图片,请检查路径是否正确")
        
        # 转化到灰度
        target = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)
        # 返回绝对值
        target = abs(255 - target)
        # 单通道转3通道
        target = cv2.cvtColor(target, cv2.COLOR_GRAY2RGB)
        template = cv2.cvtColor(template, cv2.COLOR_GRAY2RGB)
        # 进行匹配
        result = cv2.matchTemplate(target, template, cv2.TM_CCOEFF_NORMED)
        # 通过np转化为数值，就是坐标
        x, y = np.unravel_index(result.argmax(), result.shape)
        return y, x

    def from_buffer_get_distanct(self, tag_img, background_img):
        '''
        根据二进制进行识别
        :param tag_img_path: 滑块图片的二进制
        :param bg: 背景图片的二进制
        :return:
        '''
        target = cv2.imdecode(np.frombuffer(tag_img, np.uint8), cv2.IMREAD_COLOR)

        # 如果是PIL.images就换读取方式
        template = cv2.imdecode(np.frombuffer(background_img, np.uint8), cv2.IMREAD_COLOR) if type(
            background_img) == bytes else cv2.cvtColor(
            np.asarray(background_img), cv2.COLOR_RGB2BGR)

        # 转化到灰度
        target = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)

        # 返回绝对值
        target = abs(255 - target)

        # 单通道转3通道
        target = cv2.cvtColor(target, cv2.COLOR_GRAY2RGB)

        # 进行匹配
        result = cv2.matchTemplate(target, template, cv2.TM_CCOEFF_NORMED)

        # 通过np转化为数值，就是坐标
        x, y = np.unravel_index(result.argmax(), result.shape)
        return y, x

    def get_distance(self, background_img_path, tag_img_path):
        # 使用支持中文路径的方式读取图片
        def imread_chinese(filepath):
            """读取包含中文路径的图片"""
            img_array = np.fromfile(filepath, dtype=np.uint8)
            return cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        
        # 读取到背景图片的 rgb
        background_rgb = imread_chinese(background_img_path)
        # 读取到滑块图片的 rgb
        tag_rgb = imread_chinese(tag_img_path)
        
        # 检查图片是否读取成功
        if background_rgb is None or tag_rgb is None:
            raise ValueError("无法读取图片,请检查路径是否正确")
        
        # 计算结果
        res = cv2.matchTemplate(background_rgb, tag_rgb, cv2.TM_CCOEFF_NORMED)
        # 获取最小长度
        lo = cv2.minMaxLoc(res)

        # 识别返回滑动距离
        return lo[2][0]


if __name__ == '__main__':
    tag_img_path = os.path.join(os.path.dirname(__file__), 'tag.png')
    background_img_path = os.path.join(os.path.dirname(__file__), 'background.png')
    with open(tag_img_path, 'rb') as f:
        tag_img = f.read()
    with open(background_img_path, 'rb') as f:
        background_img = f.read()
    cv_obj = CvImageDistance()

    # 通过二进制数据获取缺口位置
    print(cv_obj.from_buffer_get_distanct(tag_img, background_img))
    # 通过文件位置获取缺口位置
    print(cv_obj.get_distance(background_img_path, tag_img_path))